package com.sicheng.lc.周赛.分类.离散化;

import java.util.ArrayList;

/**
 * @author zsc
 * @version 1.0
 * @date 2022/8/5 17:31
 */
public class 区间和的个数 {
    //https://leetcode.cn/problems/count-of-range-sum/

    int N = (int) 1e6;
    int[] tree = new int[N + 1];

    void insert(int x) {
        for (int i = x; i <= N; i += i & -i) {
            tree[i]++;
        }
    }

    int query(int x) {
        int sum = 0;
        for (int i = x; i > 0; i -= i & -i) {
            sum += tree[i];
        }
        return sum;
    }

    int query(int l, int r) {
        return query(r) - query(l - 1);
    }

    public int countRangeSum(int[] nums, int L, int R) {
        ArrayList<Long> t = new ArrayList<>(nums.length << 1 | 1);
        t.add(0L);
        long[] s = new long[nums.length + 1];

        for (int i = 1; i < s.length; i++) {
            s[i] += s[i - 1] + nums[i - 1];
            t.add(s[i]);
            t.add(s[i] - R);
            t.add(s[i] - L);
        }


        long[] a = t
                .stream()
                .distinct()
                .sorted()
                .mapToLong(k -> k)
                .toArray();
        int i = leftSearch(a, 0);
        insert(i);

        int res = 0;
        for (int j = 1; j < s.length; j++) {
            int l = leftSearch(a, s[j] - R);
            int r = leftSearch(a, s[j] - L);
            res += query(l, r);
            insert(leftSearch(a, s[j]));
        }

        return res;

    }


    int leftSearch(long[] nums, long x) {
        int l = 0, r = nums.length - 1;
        while (l <= r) {
            int mid = l + r >> 1;
            if (nums[mid] >= x) {
                r = mid - 1;
            } else
                l = mid + 1;
        }
        return l + 1;
    }
}
